"""Provides BUILD macros for MediaPipe graphs.

mediapipe_binary_graph() converts a graph from text format to serialized binary
format.

Example:
  mediapipe_binary_graph(
    name = "make_graph_binarypb",
    graph = "//mediapipe/framework/tool/testdata:test_graph",
    output_name = "test.binarypb",
    deps = [
        "//video/annotation:graph_calculators_lib",
    ]
  )

"""

load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
load("//mediapipe/framework:encode_binary_proto.bzl", "encode_binary_proto", "generate_proto_descriptor_set")
load("//mediapipe/framework:transitive_protos.bzl", "transitive_proto_cc_libs", "transitive_proto_descriptor_sets", "transitive_protos")
load("//mediapipe/framework/deps:expand_template.bzl", "expand_template")
load("//mediapipe/framework/tool:build_defs.bzl", "clean_dep")

# buildifier: disable=out-of-order-load
# buildifier: disable=same-origin-load
load("//mediapipe/framework/deps:descriptor_set.bzl", "direct_descriptor_set", "transitive_descriptor_set")
load("@org_tensorflow//tensorflow/lite/core/shims:cc_library_with_tflite.bzl", "cc_library_with_tflite")

def mediapipe_binary_graph(name, graph = None, output_name = None, deps = [], testonly = None, compatible_with = None, **kwargs):
    """Converts a graph from text format to binary format.

    Args:
      name: the name of the encode_binary_proto rule generated by this macro.
      graph: the BUILD label of a text-format MediaPipe graph or
          a dictionary mapping configuration to label.
      output_name: the name of the file to which the binary serialization is
          written.
      deps: the BUILD labels of dependencies that provide any additional message
          types used by the graph. The basic messages defined in calculator.proto
          are always available, but any custom types (e.g. specific calculator
          options) should be provided here. It is sufficient to provide targets
          that depend on the required protos indirectly: this macro examines the
          entire dependency tree, and does not build any dependencies except for
          the protos it finds.
      testonly: pass 1 if the graph is to be used only for tests.
      compatible_with: a list of environments the rule is compatible with.
      **kwargs: any other arguments valid for encode_binary_proto.
    """

    if not graph:
        fail("No input graph file specified.")

    if not output_name:
        fail("Must specify the output_name.")

    transitive_protos(
        name = name + "_gather_protos",
        deps = deps,
        testonly = testonly,
        compatible_with = compatible_with,
    )

    # This collects descriptor sets for tools that need them.
    transitive_proto_descriptor_sets(
        name = name + "_gather_proto_descriptor_sets",
        deps = deps,
        testonly = testonly,
        compatible_with = compatible_with,
    )

    # This collects the generated .a libraries for tools that need them.
    transitive_proto_cc_libs(
        name = name + "_gather_proto_libs",
        deps = deps,
        testonly = testonly,
        compatible_with = compatible_with,
    )

    # This generates a single descriptor set with a single invocation of the proto compiler.
    # May be faster than using the descriptor sets from proto_library.
    # We always pass at least the calculator proto since the proto compiler would fail
    # if it were passed no protos.
    generate_proto_descriptor_set(
        name = name + "_proto_descriptor_set",
        deps = [
            name + "_gather_protos",
            clean_dep("//mediapipe/framework:calculator_proto"),
        ],
        testonly = testonly,
        compatible_with = compatible_with,
    )
    encode_binary_proto(
        name = name,
        deps = [
            name + "_gather_protos",
            clean_dep("//mediapipe/framework:calculator_proto"),
            clean_dep("//mediapipe/framework/stream_handler:fixed_size_input_stream_handler_proto"),
        ],
        message_type = "mediapipe.CalculatorGraphConfig",
        input = graph,
        output = output_name,
        testonly = testonly,
        compatible_with = compatible_with,
        **kwargs
    )

def data_as_c_string(
        name,
        srcs,
        outs = None,
        testonly = None,
        compatible_with = None):
    """Encodes the data from a file as a C string literal.

    This produces a text file containing the quoted C string literal. It can be
    included directly in a C++ source file.

    Args:
      name: The name of the rule.
      srcs: A list containing a single item, the file to encode.
      outs: A list containing a single item, the name of the output text file.
            Defaults to the rule name.
      testonly: pass 1 if the graph is to be used only for tests.
      compatible_with: a list of environments the rule is compatible with.
    """
    if len(srcs) != 1:
        fail("srcs must be a single-element list")
    if outs == None:
        outs = [name]
    encode_as_c_string = clean_dep("//mediapipe/framework/tool:encode_as_c_string")
    native.genrule(
        name = name,
        srcs = srcs,
        outs = outs,
        cmd = "$(location %s) \"$<\" > \"$@\"" % encode_as_c_string,
        tools = [encode_as_c_string],
        testonly = testonly,
        compatible_with = compatible_with,
    )

def mediapipe_simple_subgraph(
        name,
        register_as,
        graph,
        deps = [],
        tflite_deps = None,
        visibility = None,
        testonly = None,
        **kwargs):
    """Defines a registered subgraph for inclusion in other graphs.

    Args:
      name: name of the subgraph target to define.
      register_as: name used to invoke this graph in supergraphs. Should be in
          CamelCase.
      graph: the BUILD label of a text-format MediaPipe graph.
      deps: any calculators or subgraphs used by this graph.
      tflite_deps: any calculators or subgraphs used by this graph that may use different TFLite implementation.
      visibility: The list of packages the subgraph should be visible to.
      testonly: pass 1 if the graph is to be used only for tests.
      **kwargs: Remaining keyword args, forwarded to cc_library.
    """
    graph_base_name = name
    mediapipe_binary_graph(
        name = name + "_graph",
        graph = graph,
        output_name = graph_base_name + ".binarypb",
        deps = deps,
        testonly = testonly,
    )
    data_as_c_string(
        name = name + "_inc",
        srcs = [graph_base_name + ".binarypb"],
        outs = [graph_base_name + ".inc"],
    )

    # cc_library for a linked mediapipe graph.
    expand_template(
        name = name + "_linked_cc",
        template = clean_dep("//mediapipe/framework/tool:simple_subgraph_template.cc"),
        out = name + "_linked.cc",
        substitutions = {
            "{{SUBGRAPH_CLASS_NAME}}": register_as,
            "{{SUBGRAPH_INC_FILE_PATH}}": native.package_name() + "/" + graph_base_name + ".inc",
        },
        testonly = testonly,
    )
    if not tflite_deps:
        cc_library(
            name = name,
            srcs = [
                name + "_linked.cc",
                graph_base_name + ".inc",
            ],
            deps = [
                clean_dep("//mediapipe/framework:calculator_framework"),
                clean_dep("//mediapipe/framework:subgraph"),
            ] + deps,
            alwayslink = 1,
            visibility = visibility,
            testonly = testonly,
            **kwargs
        )
    else:
        cc_library_with_tflite(
            name = name,
            srcs = [
                name + "_linked.cc",
                graph_base_name + ".inc",
            ],
            tflite_deps = tflite_deps,
            deps = [
                clean_dep("//mediapipe/framework:calculator_framework"),
                clean_dep("//mediapipe/framework:subgraph"),
            ] + deps,
            alwayslink = 1,
            visibility = visibility,
            testonly = testonly,
            **kwargs
        )

def mediapipe_reexport_library(
        name,
        actual,
        **kwargs):
    """Defines a cc_library that exports the headers of other libraries.

    Normally cc_library does not export the headers of its dependencies,
    and the clang "layering_check" requires clients to depend on them
    directly.  Header files can be exported by listing them in either
    cc_library's "hdrs" or "textual_hdrs" argument.  The "textual_hdrs"
    argument can also accept library targets and has the effect of
    exporting their header files and permitting client references to them.
    The result is a new library target that combines and exports the public
    interfaces of several existing library targets.

    Args:
      name: the name for the combined target.
      actual: the targets to combine and export together.
      **kwargs: Remaining keyword args, forwarded to cc_library.
    """
    cc_library(
        name = name,
        textual_hdrs = actual,
        deps = actual,
        **kwargs
    )

def mediapipe_options_library(
        name,
        proto_lib,
        deps = [],
        visibility = None,
        testonly = None,
        compatible_with = None,
        **kwargs):
    """Registers options protobuf metadata for defining options packets.

    Args:
      name: name of the options_lib target to define.
      proto_lib: the proto_library target to register.
      deps: any additional protobuf dependencies.
      visibility: The list of packages the subgraph should be visible to.
      testonly: pass 1 if the graph is to be used only for tests.
      compatible_with: a list of environments the rule is compatible with.
      **kwargs: Remaining keyword args, forwarded to cc_library.
    """

    transitive_descriptor_set(
        name = proto_lib + "_transitive",
        deps = [proto_lib],
        testonly = testonly,
        compatible_with = compatible_with,
    )
    direct_descriptor_set(
        name = proto_lib + "_direct",
        deps = [proto_lib],
        testonly = testonly,
        compatible_with = compatible_with,
    )
    data_as_c_string(
        name = name + "_inc",
        srcs = [proto_lib + "_transitive-transitive-descriptor-set.proto.bin"],
        outs = [proto_lib + "_descriptors.inc"],
        compatible_with = compatible_with,
    )
    native.genrule(
        name = name + "_type_name",
        srcs = [proto_lib + "_direct-direct-descriptor-set.proto.bin"],
        outs = [name + "_type_name.h"],
        cmd = ("$(location " + "//mediapipe/framework/tool:message_type_util" + ") " +
               ("--input_path=$(location %s) " % (proto_lib + "_direct-direct-descriptor-set.proto.bin")) +
               ("--root_type_macro_output_path=$(location %s) " % (name + "_type_name.h"))),
        tools = ["//mediapipe/framework/tool:message_type_util"],
        visibility = visibility,
        testonly = testonly,
        compatible_with = compatible_with,
    )
    expand_template(
        name = name + "_cc",
        template = clean_dep("//mediapipe/framework/tool:options_lib_template.cc"),
        out = name + ".cc",
        substitutions = {
            "{{MESSAGE_NAME_HEADER}}": native.package_name() + "/" + name + "_type_name.h",
            "{{MESSAGE_PROTO_HEADER}}": native.package_name() + "/" + proto_lib.replace("_proto", ".pb.h"),
            "{{DESCRIPTOR_INC_FILE_PATH}}": native.package_name() + "/" + proto_lib + "_descriptors.inc",
        },
        testonly = testonly,
        compatible_with = compatible_with,
    )
    cc_library(
        name = proto_lib.replace("_proto", "_options_registry"),
        srcs = [
            name + ".cc",
            proto_lib + "_descriptors.inc",
            name + "_type_name.h",
        ],
        deps = [
            clean_dep("//mediapipe/framework:calculator_framework"),
            clean_dep("//mediapipe/framework/port:advanced_proto"),
            clean_dep("//mediapipe/framework/tool:options_registry"),
            proto_lib.replace("_proto", "_cc_proto"),
        ] + deps,
        alwayslink = 1,
        visibility = visibility,
        testonly = testonly,
        features = ["-no_undefined"],
        compatible_with = compatible_with,
        **kwargs
    )
    mediapipe_reexport_library(
        name = name,
        actual = [
            proto_lib.replace("_proto", "_cc_proto"),
            proto_lib.replace("_proto", "_options_registry"),
        ],
        visibility = visibility,
        testonly = testonly,
        **kwargs
    )
